{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "thorough-brake",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy.optimize import minimize\n",
    "import pandas as pd\n",
    "from IPython.display import clear_output\n",
    "import copy\n",
    "import matplotlib.pyplot as plt\n",
    "import time\n",
    "\n",
    "from FBSM_library import FBSM, Objective_Functional\n",
    "from DM_library import DM\n",
    "from InitProblem_library import RKM_Dyn_System\n",
    "from Support_library import Define_Arguments, Obj_Functional_Iterations, Show_Control, Plot_Dynamics\n",
    "from Transition_Matrices_Bank import P_empirics_5\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad582436-6667-4062-b95f-4bfa74211d50",
   "metadata": {},
   "source": [
    "# M=1, m=5 "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "def55a72-0829-4ae0-92ec-f028d3f75ed3",
   "metadata": {},
   "source": [
    "## Set parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "723c33a7-28a2-46cc-9672-cd77007a1c06",
   "metadata": {},
   "outputs": [],
   "source": [
    "M = 1  # the number of types of ordinary agents \n",
    "\n",
    "m = 5  # the number of elements in the opinion alphabet\n",
    "n = 0.6  # the fraction of ordinary agents in the system \n",
    "\n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "# Initialize the transition matrix\n",
    "\n",
    "Transition_Matrices = {}\n",
    "\n",
    "# These transition matrices outline how ordinary agents influence each other \n",
    "\n",
    "for i in range(M):\n",
    "    for j in range(M):\n",
    "        Transition_Matrices[f'{i}-{j}'] = P_empirics_5\n",
    "\n",
    "# These transition matrices outline how stubborn agents influence ordinary agents \n",
    "        \n",
    "for i in range(M):\n",
    "    Transition_Matrices[f'{i}-{M}'] = P_empirics_5\n",
    "    \n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "# Initialize the grid\n",
    "\n",
    "Tau_0 = 0\n",
    "Tau_1 = 40\n",
    "Step = 1\n",
    "Grid = np.arange(Tau_0, Tau_1+Step, Step)\n",
    "\n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "# Define the objective functional\n",
    "\n",
    "K = 1\n",
    "\n",
    "w = K*np.array([4, 3, 2, 1, 0])  # the opinion weights (integral)\n",
    "v = np.array([4, 3, 2, 1, 0])  # the opinion weights (terminal)\n",
    "\n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "# Initialize the starting point ( y_{0} )\n",
    "\n",
    "y_0 = np.array([[0.6], \n",
    "                [0], \n",
    "                [0], \n",
    "                [0], \n",
    "                [0]])\n",
    "\n",
    "\n",
    "n_types = [0.6]\n",
    "\n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "# Create a dictionary of arguments\n",
    "\n",
    "Arguments = Define_Arguments(M, m, \n",
    "                             n, \n",
    "                             Transition_Matrices, \n",
    "                             Step, Grid, Tau_0, Tau_1, \n",
    "                             w, v, \n",
    "                             n_types)\n",
    "\n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "# Initialize the starting control guess ( u_{0} (\\tau) )\n",
    "\n",
    "u_basic = np.array([[0.08], \n",
    "                    [0.08], \n",
    "                    [0.08], \n",
    "                    [0.08], \n",
    "                    [0.08]])\n",
    "\n",
    "\n",
    "u_t_0 = []\n",
    "\n",
    "for i in range(len(Grid)-1):\n",
    "    \n",
    "    u_t_0.append(u_basic) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "299cf628-d798-43b0-b1c3-387fec09629d",
   "metadata": {},
   "source": [
    "## Run the direct method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec43e141-bfac-4bf2-beee-94b9c720cd59",
   "metadata": {},
   "outputs": [],
   "source": [
    "t_start = time.time()\n",
    "\n",
    "u_t = DM(y_0, u_t_0, Arguments)\n",
    "\n",
    "t_end = time.time()\n",
    "\n",
    "t_elapsed = round((t_end - t_start) / 60, 1)\n",
    "\n",
    "t_elapsed_sec = round((t_end - t_start), 1)\n",
    "\n",
    "print(f'Time elapsed: {t_elapsed} min / {t_elapsed_sec} sec')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3a99dfe-164b-49d9-b57f-0541c1f2a3f0",
   "metadata": {},
   "source": [
    "## Show results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b984d8bf-41fd-4706-99d9-8b11a5610769",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Show the control function obtained\n",
    "\n",
    "print('')\n",
    "\n",
    "print('The control function obtained:')\n",
    "\n",
    "print(Show_Control(u_t))\n",
    "\n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "# Calculate the value of the objective functional\n",
    "\n",
    "y_t = RKM_Dyn_System(y_0, u_t, Arguments, 1)\n",
    "\n",
    "value = Objective_Functional(y_t, Arguments)\n",
    "\n",
    "print('')\n",
    "\n",
    "print('The value of the objective functional provided by this control:', round(value[0]+value[1], 3))\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1c170f4-f6a8-48ef-8cbc-1f451b4882fd",
   "metadata": {},
   "source": [
    "## Run the FBS method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e036b3b-c011-496c-bbb3-64ff445c2a2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "t_start = time.time()\n",
    "\n",
    "(u_t_estimations, Functional_estimations, Status) = FBSM(y_0, u_t_0, Arguments)\n",
    "\n",
    "t_end = time.time()\n",
    "\n",
    "t_elapsed = round((t_end - t_start) / 60, 1)\n",
    "\n",
    "t_elapsed_sec = round((t_end - t_start), 1)\n",
    "\n",
    "print(f'Time elapsed: {t_elapsed} min / {t_elapsed_sec} sec')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2d2219e-fed6-4e68-8a62-09e91f2ba25a",
   "metadata": {},
   "source": [
    "## Show results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01a9aea6-0845-416f-92b4-58cbcfc06979",
   "metadata": {},
   "outputs": [],
   "source": [
    "print('The last approximation:')\n",
    "\n",
    "Show_Control(u_t_estimations[-1])\n",
    "\n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "print('')\n",
    "print('The best approximation:')\n",
    "\n",
    "sums = []\n",
    "\n",
    "for i in Functional_estimations:\n",
    "    sums.append(i[0]+i[1])\n",
    "\n",
    "Show_Control(u_t_estimations[np.argmin(sums)])\n",
    "\n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "y_t = RKM_Dyn_System(y_0, u_t_estimations[-1], Arguments, 1)  # last approximation\n",
    "value = Objective_Functional(y_t, Arguments)\n",
    "print('The value of the objective functional provided by the last approximation:', round(value[0]+value[1], 3))\n",
    "\n",
    "y_t = RKM_Dyn_System(y_0, u_t_estimations[np.argmin(sums)], Arguments, 1)  # the best approximation\n",
    "value = Objective_Functional(y_t, Arguments)\n",
    "print('')\n",
    "print('The value of the objective functional provided by the best approximation:', round(value[0]+value[1], 3))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a290a432-57d1-4607-ba2d-7f6118637d85",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show the dynamics of the objective functional across iterations of the FBS method\n",
    "\n",
    "FS = 15\n",
    "LW = 2\n",
    "MS = 5\n",
    "MEW = 2\n",
    "LS = 10\n",
    "\n",
    "Obj_Functional_Iterations(Functional_estimations, FS, LW, MS, MEW, LS)\n",
    "\n",
    "# Plot the corresponding dynamics of state variables (the last approximation)\n",
    "\n",
    "y_t = RKM_Dyn_System(y_0, u_t_estimations[-1], Arguments, 1)  \n",
    "\n",
    "LW = 1\n",
    "FS = 15\n",
    "\n",
    "Plot_Dynamics(y_t, LW, FS, Arguments)\n",
    "\n",
    "# Plot the corresponding dynamics of state variables (the best approximation)\n",
    "\n",
    "y_t = RKM_Dyn_System(y_0, u_t_estimations[np.argmin(sums)], Arguments, 1)  \n",
    "\n",
    "LW = 1\n",
    "FS = 15\n",
    "\n",
    "Plot_Dynamics(y_t, LW, FS, Arguments)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "660f9d8f-8f0d-4c34-afb0-4e006195751b",
   "metadata": {},
   "source": [
    "## Compare against benchmark controls: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f533bad-4bed-4044-a737-e5667863c4f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "u_basic = np.array([[0], \n",
    "                    [0], \n",
    "                    [0], \n",
    "                    [0.4], \n",
    "                    [0]])\n",
    "\n",
    "\n",
    "u_t = []\n",
    "\n",
    "for i in range(len(Grid)-1):\n",
    "    \n",
    "    u_t.append(u_basic) \n",
    "    \n",
    "y_t = RKM_Dyn_System(y_0, u_t, Arguments, 1)\n",
    "value = Objective_Functional(y_t, Arguments)\n",
    "print(round(value[0]+value[1], 5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f60259f5-8d0c-47ca-b77c-791d11f88ae4",
   "metadata": {},
   "outputs": [],
   "source": [
    "u_basic = np.array([[0], \n",
    "                    [0], \n",
    "                    [0], \n",
    "                    [0], \n",
    "                    [0.4]])\n",
    "\n",
    "\n",
    "u_t = []\n",
    "\n",
    "for i in range(len(Grid)-1):\n",
    "    \n",
    "    u_t.append(u_basic) \n",
    "    \n",
    "y_t = RKM_Dyn_System(y_0, u_t, Arguments, 1)\n",
    "value = Objective_Functional(y_t, Arguments)\n",
    "print(round(value[0]+value[1], 5))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a8150af-6d16-444c-bdf9-32691fb10c63",
   "metadata": {},
   "source": [
    "# Try a different objective functional - pushing the system towards the left"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b395404f-e5cb-45f5-9e9e-483163af18d8",
   "metadata": {},
   "source": [
    "## Set parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5fa2866-5881-49c5-a250-25718bf74892",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Define the objective functional\n",
    "\n",
    "K = 1\n",
    "\n",
    "v = np.array([0, 1, 2, 3, 4])  # the opinion weights (terminal)\n",
    "w = K*v  # the opinion weights (integral)\n",
    "\n",
    "\n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "# Create a dictionary of arguments\n",
    "\n",
    "Arguments = Define_Arguments(M, m, \n",
    "                             n, \n",
    "                             Transition_Matrices, \n",
    "                             Step, Grid, Tau_0, Tau_1, \n",
    "                             w, v, \n",
    "                             n_types)\n",
    "\n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "# Initialize the starting point ( y_{0} )\n",
    "\n",
    "y_0 = np.array([[0], \n",
    "                [0], \n",
    "                [0], \n",
    "                [0], \n",
    "                [0.6]])\n",
    "\n",
    "\n",
    "n_types = [0.6]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f528a3a0-4a54-4e0c-9bfd-1ea81e7fca39",
   "metadata": {},
   "source": [
    "## Run the direct method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a15e40b-428c-41d6-b272-7fe6298d8fdd",
   "metadata": {},
   "outputs": [],
   "source": [
    "t_start = time.time()\n",
    "\n",
    "u_t = DM(y_0, u_t_0, Arguments)\n",
    "\n",
    "t_end = time.time()\n",
    "\n",
    "t_elapsed = round((t_end - t_start) / 60, 1)\n",
    "\n",
    "t_elapsed_sec = round((t_end - t_start), 1)\n",
    "\n",
    "print(f'Time elapsed: {t_elapsed} min / {t_elapsed_sec} sec')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55791aa8-07d6-4a60-9baa-e90061aa2eaf",
   "metadata": {},
   "source": [
    "## Show results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3336219-2de4-4568-9294-c307ec7822f6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Show the control function obtained\n",
    "\n",
    "print('')\n",
    "\n",
    "print('The control function obtained:')\n",
    "\n",
    "print(Show_Control(u_t))\n",
    "\n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "# Calculate the value of the objective functional\n",
    "\n",
    "y_t = RKM_Dyn_System(y_0, u_t, Arguments, 1)\n",
    "\n",
    "value = Objective_Functional(y_t, Arguments)\n",
    "\n",
    "print('')\n",
    "\n",
    "print('The value of the objective functional provided by this control:', round(value[0]+value[1], 3))\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4826c58-c5eb-4113-be00-95816d6cfa73",
   "metadata": {},
   "source": [
    "## Run the FBS method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "832444cc-e2b6-4766-8a57-5dfa6f5c4a75",
   "metadata": {},
   "outputs": [],
   "source": [
    "t_start = time.time()\n",
    "\n",
    "(u_t_estimations, Functional_estimations, Status) = FBSM(y_0, u_t_0, Arguments)\n",
    "\n",
    "t_end = time.time()\n",
    "\n",
    "t_elapsed = round((t_end - t_start) / 60, 1)\n",
    "\n",
    "t_elapsed_sec = round((t_end - t_start), 1)\n",
    "\n",
    "print(f'Time elapsed: {t_elapsed} min / {t_elapsed_sec} sec')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6bdcce7e-97a9-491c-854c-d5e15b39a78c",
   "metadata": {},
   "source": [
    "## Show results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5c71716-dab7-47dd-a85e-6ee2602ebf31",
   "metadata": {},
   "outputs": [],
   "source": [
    "print('The last approximation:')\n",
    "\n",
    "Show_Control(u_t_estimations[-1])\n",
    "\n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "print('')\n",
    "print('The best approximation:')\n",
    "\n",
    "sums = []\n",
    "\n",
    "for i in Functional_estimations:\n",
    "    sums.append(i[0]+i[1])\n",
    "\n",
    "Show_Control(u_t_estimations[np.argmin(sums)])\n",
    "\n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "y_t = RKM_Dyn_System(y_0, u_t_estimations[-1], Arguments, 1)  # last approximation\n",
    "value = Objective_Functional(y_t, Arguments)\n",
    "print('The value of the objective functional provided by the last approximation:', round(value[0]+value[1], 3))\n",
    "\n",
    "y_t = RKM_Dyn_System(y_0, u_t_estimations[np.argmin(sums)], Arguments, 1)  # the best approximation\n",
    "value = Objective_Functional(y_t, Arguments)\n",
    "print('')\n",
    "print('The value of the objective functional provided by the best approximation:', round(value[0]+value[1], 3))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7bb2e980-aad3-4c59-83c6-df84673a81e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show the dynamics of the objective functional across iterations of the FBS method\n",
    "\n",
    "FS = 15\n",
    "LW = 2\n",
    "MS = 5\n",
    "MEW = 2\n",
    "LS = 10\n",
    "\n",
    "Obj_Functional_Iterations(Functional_estimations, FS, LW, MS, MEW, LS)\n",
    "\n",
    "# Plot the corresponding dynamics of state variables (the last approximation)\n",
    "\n",
    "y_t = RKM_Dyn_System(y_0, u_t_estimations[-1], Arguments, 1)  \n",
    "\n",
    "LW = 1\n",
    "FS = 15\n",
    "\n",
    "Plot_Dynamics(y_t, LW, FS, Arguments)\n",
    "\n",
    "# Plot the corresponding dynamics of state variables (the best approximation)\n",
    "\n",
    "y_t = RKM_Dyn_System(y_0, u_t_estimations[np.argmin(sums)], Arguments, 1)  \n",
    "\n",
    "LW = 1\n",
    "FS = 15\n",
    "\n",
    "Plot_Dynamics(y_t, LW, FS, Arguments)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52b4a2f0-3a7c-4cd6-b03a-94494e99425a",
   "metadata": {},
   "source": [
    "## Compare against benchmark controls: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "415bd20b-8b13-4719-ba5f-e49af6b0dc18",
   "metadata": {},
   "outputs": [],
   "source": [
    "u_basic = np.array([[0.4], \n",
    "                    [0], \n",
    "                    [0], \n",
    "                    [0], \n",
    "                    [0]])\n",
    "\n",
    "\n",
    "u_t = []\n",
    "\n",
    "for i in range(len(Grid)-1):\n",
    "    \n",
    "    u_t.append(u_basic) \n",
    "    \n",
    "y_t = RKM_Dyn_System(y_0, u_t, Arguments, 1)\n",
    "value = Objective_Functional(y_t, Arguments)\n",
    "print(round(value[0]+value[1], 5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3bc2b461-8350-4b8f-aa6f-5024f4fde3a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "u_basic = np.array([[0], \n",
    "                    [0.4], \n",
    "                    [0], \n",
    "                    [0], \n",
    "                    [0]])\n",
    "\n",
    "\n",
    "u_t = []\n",
    "\n",
    "for i in range(len(Grid)-1):\n",
    "    \n",
    "    u_t.append(u_basic) \n",
    "    \n",
    "y_t = RKM_Dyn_System(y_0, u_t, Arguments, 1)\n",
    "value = Objective_Functional(y_t, Arguments)\n",
    "print(round(value[0]+value[1], 5))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
